module aru_arb_wrgen_addr_calc (
    input logic                             clk,
    input logic                             rst_n,
          aru_arb_wrgen_cfg_if.addr_calc_in u_aru_cfg_if,
          aru_payload_if.in                 u_aru_pld_if,
          aru_idx_if.in                     u_aru_idx_if,
          aru_sdb_if.in                     u_aru_sdb_if,
          arb_wr_req_if.out                 u_arb_wr_req_if,
          done_if.out                       u_done_if
);

    logic cfg_vld;
    logic lst_req_in_instr;
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            cfg_vld <= 1'b0;
        end else if (cfg_vld) begin
            if (u_arb_wr_req_if.vld && u_arb_wr_req_if.rdy && lst_req_in_instr) begin
                u_done_if.vld       <= 1'b1;
                u_done_if.instr_idx <= u_aru_cfg_if.instr_idx;
                cfg_vld             <= 1'b0;
            end
        end else begin
            if (u_aru_cfg_if.vld) begin
                cfg_vld <= 1'b1;
            end
        end
    end
    assign u_aru_cfg_if.rdy = !cfg_vld;
    assign lst_req_in_instr = u_aru_sdb_if.pld.eom && u_aru_sdb_if.pld.eon;

    idx_t slice_m1;
    idx_t slice_n1;
    logic msk      [`P_ARU-1:0][`N0-1:0];

    assign slice_m1 = (u_aru_cfg_if.slice_m + `M0 - 1) / `M0;
    assign slice_n1 = (u_aru_cfg_if.slice_n + `N0 - 1) / `N0;

    // Generate mask based on valid dimensions
    always_comb begin
        for (integer i = 0; i < `P_ARU; i = i + 1) begin
            for (integer j = 0; j < `N0; j = j + 1) begin
                if (i < u_aru_sdb_if.pld.vld_m && j < u_aru_sdb_if.pld.vld_n) begin
                    msk[i][j] = 1'b1;
                end else begin
                    msk[i][j] = 1'b0;
                end
            end
        end
    end

    // Address and data calculation
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            u_arb_wr_req_if.arb_addr <= 'd0;
            u_arb_wr_req_if.dat      <= 'd0;
            u_arb_wr_req_if.reduce_m <= 'd0;
            u_arb_wr_req_if.reduce_n <= 'd0;
            u_arb_wr_req_if.msk      <= 'd0;
        end else begin
            // Address calculation based on reduce mode
            if (u_aru_cfg_if.reduce_m && u_aru_cfg_if.reduce_n) begin
                // Reduce MN mode - single output
                u_arb_wr_req_if.arb_addr <= u_aru_cfg_if.arb_addr;
                u_arb_wr_req_if.reduce_m <= 1'b1;
                u_arb_wr_req_if.reduce_n <= 1'b1;
            end else if (u_aru_cfg_if.reduce_m) begin
                // Reduce M mode - increment by 16B
                u_arb_wr_req_if.arb_addr <= u_aru_cfg_if.arb_addr + (u_aru_idx_if.pld.n1_idx * 2);
                u_arb_wr_req_if.reduce_m <= 1'b1;
                u_arb_wr_req_if.reduce_n <= 1'b0;
            end else if (u_aru_cfg_if.reduce_n) begin
                // Reduce N mode - increment by 8B
                u_arb_wr_req_if.arb_addr <= u_aru_cfg_if.arb_addr + 
                                          (u_aru_idx_if.pld.m1_idx * `M0 + u_aru_idx_if.pld.p_idx * `P_ARU) * 2;
                u_arb_wr_req_if.reduce_m <= 1'b0;
                u_arb_wr_req_if.reduce_n <= 1'b1;
            end else begin
                // Normal mode - increment by 64B
                u_arb_wr_req_if.arb_addr <= u_aru_cfg_if.arb_addr + 
                                          (u_aru_idx_if.pld.n1_idx * slice_m1 * `M0 + 
                                          (u_aru_idx_if.pld.m1_idx * `M0 + u_aru_idx_if.pld.p_idx * `P_ARU)) * 2;
                u_arb_wr_req_if.reduce_m <= 1'b0;
                u_arb_wr_req_if.reduce_n <= 1'b0;
            end

            // Data assignment - always pack to lower bits for reduce modes
            for (int i = 0; i < `P_ARU; i++) begin
                for (int j = 0; j < `N0; j++) begin
                    u_arb_wr_req_if.dat[i][j] <= u_aru_pld_if.dat[i*`N0+j];
                    u_arb_wr_req_if.msk[i][j] <= msk[i][j];
                end
            end
        end
    end

    // Flow control logic
    logic has_data;
    logic up_vld, up_rdy, down_vld, down_rdy;
    assign down_rdy = u_arb_wr_req_if.rdy;
    assign up_vld   = u_aru_idx_if.vld && u_aru_sdb_if.vld && u_aru_pld_if.vld;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            has_data <= 'd0;
        end else if (has_data) begin
            if (~up_vld && down_rdy) begin
                has_data <= 'd0;
            end
        end else begin
            if (up_vld) begin
                has_data <= 'd1;
            end
        end
    end

    // Output assignments
    assign down_vld            = has_data;
    assign up_rdy              = (~has_data) || down_rdy;
    assign u_arb_wr_req_if.vld = cfg_vld && down_vld;
    assign u_aru_idx_if.rdy    = up_rdy && cfg_vld;
    assign u_aru_sdb_if.rdy    = up_rdy && cfg_vld;
    assign u_aru_pld_if.rdy    = up_rdy && cfg_vld;

endmodule
